""" The defination and basic methods of bbox
"""
import numpy as np
from copy import deepcopy
import bisect
from .bbox import BBox

class BBoxSeq:
    def __init__(self, bbox: BBox, stamp):
        self.stamps = []
        self.stamps.append(stamp)
        self.lx = []
        self.lx.append(bbox.x)
        self.ly = []
        self.ly.append(bbox.y)
        self.lz = []
        self.lz.append(bbox.z)
        self.lo = []
        self.lo.append(bbox.o)
        self.lvx = []
        self.lvx.append(bbox.vx)
        self.lvy = []
        self.lvy.append(bbox.vy)
        self.lvz = []
        self.lvz.append(bbox.vz)
        self.s = bbox.s
        self.l = bbox.l
        self.h = bbox.h
        self.w = bbox.w

    def insert(self, bbox: BBox, stamp):
        self.stamps.append(stamp)
        self.lx.append(bbox.x)
        self.ly.append(bbox.y)
        self.lz.append(bbox.z)
        self.lo.append(bbox.o)
        self.lvx.append(bbox.vx)
        self.lvy.append(bbox.vy)
        self.lvz.append(bbox.vz)

    def _predict_future(self, stamp):
        """未来时间预测（使用最后速度）"""
        delta_t = (stamp - self.stamps[-1])/1000.0 # ms -> s
        return BBox(
            x=self.lx[-1] + self.lvx[-1] * delta_t,
            y=self.ly[-1] + self.lvy[-1] * delta_t,
            z=self.lz[-1] + self.lvz[-1] * delta_t,
            o=self.lo[-1],  # 方向保持最后记录值
            vx=self.lvx[-1],
            vy=self.lvy[-1],
            vz=self.lvz[-1],
            s=self.s,
            l=self.l,
            h=self.h,
            w=self.w
        )

    def _create_bbox(self, index):
        """创建指定索引的BBox"""
        return BBox(
            x=self.lx[index],
            y=self.ly[index],
            z=self.lz[index],
            o=self.lo[index],
            vx=self.lvx[index],
            vy=self.lvy[index],
            vz=self.lvz[index],
            s=self.s,
            l=self.l,
            h=self.h,
            w=self.w
        )

    def _interpolate(self, values, index, ratio):
        """线性插值通用方法"""
        return values[index - 1] + (values[index] - values[index - 1]) * ratio

    def _interpolate_angle(self, angles, index, ratio):
        """角度插值（考虑周期性）"""
        a0, a1 = angles[index - 1], angles[index]
        delta = (a1 - a0 + np.pi) % (2 * np.pi) - np.pi
        return (a0 + delta * ratio + np.pi) % (2 * np.pi) - np.pi
    def search(self, stamp):
        # 处理未来时间预测
        if stamp >= self.stamps[-1]:
            return self._predict_future(stamp)
        # 处理历史时间插值
        index = bisect.bisect_left(self.stamps, stamp)
        # 处理早于所有记录时间的情况
        if index == 0:
            return self._create_bbox(0)
        # 正常插值逻辑
        t0, t1 = self.stamps[index - 1], self.stamps[index]
        ratio = (stamp - t0) / (t1 - t0) if t0 != t1 else 0.0

        return BBox(
            x=self._interpolate(self.lx, index, ratio),
            y=self._interpolate(self.ly, index, ratio),
            z=self._interpolate(self.lz, index, ratio),
            o=self._interpolate_angle(self.lo, index, ratio),
            vx=self._interpolate(self.lvx, index, ratio),
            vy=self._interpolate(self.lvy, index, ratio),
            vz=self._interpolate(self.lvz, index, ratio),
            s=self.s,
            l=self.l,
            h=self.h,
            w=self.w
        )


class SearchBBox:
    def __init__(self):
        self.map = dict()

    def insert(self, id, bbox: BBox, stamp):
        if id not in self.map.keys():
            self.map[id] = BBoxSeq(bbox, stamp)
        else:
            self.map[id].insert(bbox, stamp)

    def search(self, id, time):
        if id not in self.map:
            raise KeyError(f"Target ID {id} not tracked")
        return self.map[id].search(time)

